from torch.utils.data import DataLoader
from torchvision import datasets,transforms
from configs import TEST_BATCH_SIZE, TRAIN_BATCH_SIZE

pipeline = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,),(0.3081,))
])

def prepare_train_data():
    data_train = datasets.MNIST(root = 'data/',transform = pipeline,train = True,download=True)
    train_loader = DataLoader(dataset = data_train,shuffle=True,batch_size=TRAIN_BATCH_SIZE)
    return train_loader

def prepare_test_data():
    data_test = datasets.MNIST(root='data',transform = pipeline,train=False,download = True)
    test_loader = DataLoader(data_test,batch_size=TEST_BATCH_SIZE,shuffle=True)
    return test_loader